import functools
import re

from langchain.agents import AgentExecutor, create_react_agent, create_json_chat_agent, create_structured_chat_agent
from langchain.memory import ConversationBufferMemory
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_openai import OpenAI, ChatOpenAI
from langchain.prompts import PromptTemplate, ChatPromptTemplate, SystemMessagePromptTemplate
from langchain.tools import BaseTool, StructuredTool, tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import AIMessage, HumanMessage
from langchain.agents.format_scratchpad.openai_tools import (
    format_to_openai_tool_messages,
)
from langchain_core.messages import SystemMessage,AIMessage,ChatMessage
from lc_tools import *
from langchain.agents.output_parsers import JSONAgentOutputParser

llm = ChatOpenAI(model_name="Qwen1.5-72b-chat-gptq-int4", openai_api_base="http://113.31.110.212:9987/v1",
                 openai_api_key="<KEY>")


system = """
Respond to the human as helpfully and accurately as possible. You should attempt to understand the input before you answer. You have access to the following tools:

{tools}

Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).

Valid "action" values: "Final Answer" or {tool_names}

Provide only ONE action per $JSON_BLOB, as shown:

```
{{
  "action": $TOOL_NAME,
  "action_input": $INPUT
}}
```

Follow this format:

Question: input question to answer
Thought: consider if the previous answer and the question.
 if not, expand it.   
Action:
```
$JSON_BLOB
```
Observation: action result
... (repeat Thought/Action/Observation N times)
Thought: I know what to respond
Action:
```
{{
  "action": "Final Answer",
  "action_input": "Final response to human"
}}

Previous chat history consisting of messages between human and AI:

{chat_history}


Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools to . Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation

"""

human = """
Human:
{input}

{agent_scratchpad}
 (reminder to respond in a JSON blob no matter what)
"""

tools = [count_word, get_word_length]
prompts = ChatPromptTemplate.from_messages([
    ("system", system),
    ("human", human)

])
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

write_agent = create_structured_chat_agent(llm, tools, prompts)

write_agent_executor = AgentExecutor(
    agent=write_agent, tools=tools, verbose=True, handle_parsing_errors=True,memory=memory
)



divider_template = """
我要写一篇论文关于 {input} 内容为{content},你帮我进行段落的划分

输出格式,需要是JSON的格式，至多划分只三级标题，如下是一个输出的例子


{{
    "1.一级标题":{{
        "1.1 二级标题":[
            "1.1.1 三级标题",
            ....
        ],
        "1.2 二级标题":[
            ...
        ]   
    }}
    ...

}}
"""


segment_divider_prompt = ChatPromptTemplate.from_template(divider_template)

section_divider_agent = segment_divider_prompt | llm


check_prompt = """
Respond to the human as helpfully and accurately as possible.
 You should attempt to understand the input before you answer. You have access to the following tools:

{tools}

Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).

Valid "action" values: "Final Answer" or {tool_names}

Provide only ONE action per $JSON_BLOB, as shown:

```
{{
  "action": $TOOL_NAME,
  "action_input": $INPUT
}}
```

Follow this format:

Question: input question to answer
Thought: consider if the previous answer and the question.
 if not, expand it.   
Action:
```
$JSON_BLOB
```
Observation: action result
... (repeat Thought/Action/Observation N times)
Thought: I know what to respond
Action:
```
{{
  "action": "Final Answer",
  "action_input": "number only"
}}


"""

check_input = """
Human:
how many number of words in paragraph  " {input} " 


{agent_scratchpad}
 (reminder to respond in a JSON blob no matter what)
"""



check_agent_prompt = ChatPromptTemplate.from_messages([
    ("system", check_prompt),
    ("human", check_input)
])

check_agent = create_structured_chat_agent(llm, tools, check_agent_prompt)

check_agent_exec = AgentExecutor(
    agent=check_agent, tools=tools, verbose=True, handle_parsing_errors=True,
)


writer_prompt = """
你需要根据我给出的话题进行写作,我给出的话题是,请详细的进行写作，不需要分点，用一段话即可。


输出格式，不需要输出标题，请直接输出内容，并且不要输出其他不相关的东西。

"""
writer =  ChatOpenAI(model_name="Qwen1.5-72b-chat-gptq-int4", openai_api_base="http://113.31.110.212:9987/v1",
                 openai_api_key="<KEY>")

write_history_message = [
    SystemMessage(content=writer_prompt),
]




if __name__ == '__main__':

    message = writer.invoke([*write_history_message,
                             HumanMessage(
                                 content=f"我的目标是写一些篇 基于Vus.js的美食分享管理系统的开发与实践。，内容包括美食的收藏、分享、浏览评论功能、管理员：用户管理、美食管理、评论管理, 请你帮我撰写1.1介绍SpringBoot")
                             ])

    print(message)